package ssh

import (
	"context"
	"io"
	"log"
	"os"
	"os/exec"
	"runtime"

	"golang.org/x/crypto/ssh"
)

type SessReqMsgEnv struct {
	Name  string
	Value string
}

type SessReqMsgExec struct {
	Command string
}

type SessReqMsgExitStatus struct {
	Status uint32
}

func SessSendExitStatus(ch ssh.Channel, status int) {
	ch.SendRequest("exit-status", false, ssh.Marshal(&SessReqMsgExitStatus{
		Status: uint32(status),
	}))
}

type SessionHandler struct {
	CheckAuth   func() error
	WatchConfig WatchConfigFunc
}

func (handler *SessionHandler) Handle(newCh ssh.NewChannel) {
	if newCh.ChannelType() != "session" {
		log.Panicf("ssh: unknown channel type %v", newCh.ChannelType())
	}
	if err := handler.CheckAuth(); err != nil {
		newCh.Reject(ssh.Prohibited, err.Error())
		return
	}
	ch, reqs, err := newCh.Accept()
	if err != nil {
		return
	}
	defer ch.Close()
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	go func() {
		watcher := handler.WatchConfig()
		defer watcher.Unwatch()
	Loop:
		for {
			select {
			case <-ctx.Done():
				break Loop
			case _, ok := <-watcher.OnChange:
				if !ok {
					break Loop
				}
			}
			if err := handler.CheckAuth(); err != nil {
				ch.Close()
			}
		}
	}()
	var cmd *exec.Cmd
	for req := range reqs {
		Reply(req, func() (bool, []byte) {
			switch req.Type {
			case "exec":
				if cmd != nil {
					return false, nil
				}
				var msg SessReqMsgExec
				if err := ssh.Unmarshal(req.Payload, &msg); err != nil {
					return false, nil
				}
				switch runtime.GOOS {
				case "windows":
					cmd = exec.CommandContext(ctx, "powershell", "-c", msg.Command)
				case "linux":
					cmd = exec.CommandContext(ctx, "sh", "-c", msg.Command)
				default:
					return false, nil
				}
				cmd.Env = os.Environ()
				stdin, err := cmd.StdinPipe()
				if err != nil {
					return false, nil
				}
				stdout, err := cmd.StdoutPipe()
				if err != nil {
					return false, nil
				}
				stderr, err := cmd.StderrPipe()
				if err != nil {
					return false, nil
				}
				if err := cmd.Start(); err != nil {
					return false, nil
				}
				go func() {
					io.Copy(ch, stdout)
					// ch.CloseWriteStdout()
					stdout.Close()
				}()
				go func() {
					io.Copy(ch.Stderr(), stderr)
					// ch.CloseWriteStderr()
					stderr.Close()
				}()
				go func() {
					io.Copy(stdin, ch)
					stdin.Close()
					// ch.CloseRead()
				}()
				go func() {
					cmd.Wait()
					SessSendExitStatus(ch, cmd.ProcessState.ExitCode())
					ch.Close()
				}()
				return true, nil
			default:
				return false, nil
			}
		})
	}
}
